package com.aiot.handler;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import javax.annotation.PreDestroy;
import java.util.*;
import java.util.concurrent.*;

/**
 * 精简版 VideoWebSocketHandler
 * - 根据连接路径判定 sender/viewer（/video/sender 或 /video/viewer）
 * - 发送端直接发送完整二进制帧
 * - 服务器立即广播二进制帧给所有 viewer，并更新 lastFrameBytes
 */
@Component
public class VideoWebSocketHandler extends AbstractWebSocketHandler {

    private final ConcurrentMap<String, WebSocketSession> viewers = new ConcurrentHashMap<>();
    private final ConcurrentMap<String, WebSocketSession> senders = new ConcurrentHashMap<>();

    // 最近一帧字节数据
    private volatile byte[] lastFrameBytes = null;

    // 线程池用于广播，避免阻塞 IO 线程
    private final ExecutorService broadcastPool = Executors.newFixedThreadPool(
            Math.max(2, Runtime.getRuntime().availableProcessors() / 2),
            r -> {
                Thread t = new Thread(r, "ws-broadcast-" + UUID.randomUUID());
                t.setDaemon(true);
                return t;
            }
    );

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        System.out.println("[WS] 建立连接: " + session.getId() + " uri=" + session.getUri());

        // 根据连接路径判定角色
        try {
            if (session.getUri() != null) {
                String path = session.getUri().getPath();
                if (path != null) {
                    if (path.endsWith("/sender")) {
                        senders.put(session.getId(), session);
                        System.out.println("[WS] 会话标记为发送端sender: " + session.getId());
                    } else if (path.endsWith("/viewer")) {
                        viewers.put(session.getId(), session);
                        System.out.println("[WS] 会话标记为播放端viewer: " + session.getId());
                        // 新 viewer 立即收到最近一帧（如果存在）
                        if (lastFrameBytes != null && session.isOpen()) {
                            broadcastPool.submit(() -> {
                                try {
                                    if (session.isOpen()) {
                                        session.sendMessage(new BinaryMessage(lastFrameBytes));
                                        System.out.println("[WS] 已将最近一帧发送给 viewer " + session.getId() + " 大小=" + lastFrameBytes.length + " 字节");
                                    }
                                } catch (Throwable ioe) {
                                    System.err.println("[WS] 发送最近一帧给 viewer " + session.getId() + " 失败 - " + ioe.getMessage());
                                    ioe.printStackTrace();
                                }
                            });
                        }
                    } else {
                        System.out.println("[WS] 未知路径, 会话未被标记: " + session.getId() + " path=" + path);
                    }
                }
            }
        } catch (Throwable t) {
            System.err.println("[WS] 在判定角色时发生错误: " + t.getMessage());
            t.printStackTrace();
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        // 文本消息直接忽略
        System.out.println("[WS] 收到文本消息(已忽略): 会话=" + session.getId() + " 内容=" + message.getPayload());
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
        // 仅允许已标记为 sender 的 session 发送二进制帧
        if (!senders.containsKey(session.getId())) {
            System.out.println("[WS] 收到非发送端的二进制数据, 已忽略: " + session.getId());
            return;
        }

        // 读取二进制内容到 byte[]
        byte[] bytes;
        try {
            java.nio.ByteBuffer buf = message.getPayload();
            bytes = new byte[buf.remaining()];
            buf.get(bytes);
        } catch (Throwable t) {
            System.err.println("[WS] 读取二进制数据时发生错误: " + t.getMessage());
            t.printStackTrace();
            sendError(session, "读取二进制数据错误");
            return;
        }

        // 更新最近一帧数据
        lastFrameBytes = bytes;

        // 异步广播给所有 viewers
        broadcastPool.submit(() -> {
            try {
                BinaryMessage out = new BinaryMessage(bytes);
                for (WebSocketSession viewer : viewers.values()) {
                    try {
                        if (viewer != null && viewer.isOpen()) {
                            try {
                                viewer.sendMessage(out);
                            } catch (Throwable ioe) {
                                // viewer.sendMessage 内部异常
                            }
                        }
                    } catch (Throwable ioe) {
                        System.err.println("[WS] 向观众发送数据时出错: " + (viewer != null ? viewer.getId() : "null")
                                + " - " + ioe.getMessage());
                        ioe.printStackTrace();
                    }
                }
            } catch (Throwable t) {
                System.err.println("[WS] 广播任务执行失败: " + t.getMessage());
                t.printStackTrace();
            }
        });
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        senders.remove(session.getId());
        viewers.remove(session.getId());
        System.out.println("[WS] 连接关闭: " + session.getId() + " 状态=" + status);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        System.err.println("[WS] 传输过程中发生错误: " + (session != null ? session.getId() : "null") + " - " + exception.getMessage());
        exception.printStackTrace();
    }

    private void sendError(WebSocketSession session, String reason) {
        try {
            if (session != null && session.isOpen()) {
                session.sendMessage(new TextMessage("{\"type\":\"error\",\"reason\":\"" + reason + "\"}"));
            }
        } catch (Throwable t) {
            System.err.println("[WS] 发送错误消息失败: " + t.getMessage());
            t.printStackTrace();
        }
    }

    @PreDestroy
    public void shutdown() {
        try {
            broadcastPool.shutdownNow();
        } catch (Throwable t) {
            // 忽略
        }
    }
}
